#!/usr/local/bin/python
# Fuck MakeTable5's symbols. They always make duplicate errors when working with multiple families.
# Stop using them. You dont need symbols as your 'ID'. Thats why we have gi numbers and shit. Holy shit.

import sys
import re
from Bio import SeqIO

class symbols:
	
	def __init__(self):
		self.input = None
		self.fasta = None
		self.output = None
		self.symbols = []
		self.newfasta = []
		
	def __call__(self):
		self.fasta = SeqIO.parse(self.input,'fasta')
		for fasta in self.fasta:
			new = self.assign_symbol(fasta.id)
			fasta.id = new
			self.newfasta.append(fasta)
		self.write_new()
		
	def assign_symbol(self,symbol):
		if symbol not in self.symbols:
			self.symbols.append(symbol)
			return symbol
		# shit this symbols in use. Add one. then recurse this shit..
		pattern = r'([A-Za-z]+)(\d+)'
		segments = re.search(pattern,symbol)
		sym,dig=segments.groups()[0:]
		dig = int(dig) + 1
		symbol = sym+str(dig)
		return self.assign_symbol(symbol)
		
	def write_new(self):
		handle = open(self.output,'wb')
		for entry in self.newfasta:
			SeqIO.write(entry,handle,'fasta')
		
if __name__=='__main__':
	usage = 'newsymbols.py <fasta_file> <outfile>'
	try:
		infile,outfile = sys.argv[1:3]
	except:
		print usage
		quit()
	sym=symbols()
	sym.input=infile
	sym.output=outfile
	sym()